import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

# 设置字体为 Times New Roman
plt.rcParams['font.family'] = 'Times New Roman'

# 定义拟合函数
def fit_function(t, a, b, c):
    T = t  # t 作为输入的变量
    return a * np.power(T, 2) + b * T + c

# 读取 Excel 文件中的数据
file_path = './assets/normalized_data_no_panimic.xlsx'  # Excel 文件路径
data = pd.read_excel(file_path)

N1 = data['N1'].values  # 特征 N1
T = data['T'].values  # 特征 T

# 将特征数据组合成一个元组 t，包含 T
t = T  # 转置为 (N, 1) 的二维数组

# 使用 curve_fit 进行拟合
popt, pcov = curve_fit(fit_function, t, N1, p0=[1, 1, -1])  # 初始参数猜测值

# 提取拟合参数
a, b, c = popt
print(f'拟合参数: a={a:.4f}, b={b:.4f}, c={c:.4f}')

# 计算拟合值
N1_fit = fit_function(t, a, b, c)

# 计算R方
def calculate_r_squared(y_true, y_pred):
    ss_res = np.sum((y_true - y_pred) ** 2)
    ss_tot = np.sum((y_true - np.mean(y_true)) ** 2)
    return 1 - (ss_res / ss_tot)

r_squared = calculate_r_squared(N1, N1_fit)
print(f'R²值: {r_squared:.4f}')

# 使用 seaborn 绘图
sns.set_theme(style="darkgrid")
plt.figure(figsize=(10, 6))
sns.scatterplot(x=T, y=N1, color=sns.color_palette("PuBu")[2], label='Actual Data')
sns.lineplot(x=T, y=N1_fit, color=sns.color_palette("PuBu")[4], label='Fitted Curve')
plt.ylabel('N(t)', fontsize=14, fontname='Times New Roman')
plt.title('N(t) Model Fit', fontsize=16, fontname='Times New Roman')
plt.legend()
plt.xticks(ticks=np.arange(T.min(), T.max() + 1, 1), fontname='Times New Roman')  # 设置横轴刻度为整数
plt.yticks(fontname='Times New Roman')
plt.show()